import os
import sys
import io
import time

import paramiko
import json

# 获取建木参数
# ssh连接参数
ssh_ip = os.getenv("JIANMU_SSH_IP")
ssh_port = os.getenv("JIANMU_SSH_PORT")
quantize_level = os.getenv("JIANMU_QUANTIZE_LEVEL")

if not quantize_level:
    quantize_level = 'FP16'

try:
    ssh_port = int(ssh_port)
except Exception as e:
    print("ssh端口获取错误：", e, flush=True)
    sys.exit(1)

ssh_user = os.getenv("JIANMU_SSH_USER")
ssh_pass = os.getenv("JIANMU_SSH_PASS")
ssh_private_key = os.getenv("JIANMU_SSH_PRIVATE_KEY")

# 启动参数
glm_conc_cnt = os.getenv("JIANMU_GLM_CONC_CNT")
glm_web_port = os.getenv("JIANMU_GLM_WEB_PORT")

chatglm_image = "docker.jianmuhub.com/autoops/chatglm-6b:latest"
driver_check_cmd = "nvidia-smi"
docker_path = "docker"
container_name = "autoops-chatglm-6b"
timeout_seconds = 300


def getSSHClient():
    """获取ssh client"""
    ssh_client = paramiko.SSHClient()
    ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    try:
        if ssh_private_key:
            pkey = paramiko.RSAKey.from_private_key(io.StringIO(ssh_private_key))
            ssh_client.connect(hostname=ssh_ip, port=ssh_port, username=ssh_user, pkey=pkey)
            print("Get ssh_client by pkey.", flush=True)
        elif ssh_pass:
            ssh_client.connect(hostname=ssh_ip, port=ssh_port, username=ssh_user, password=ssh_pass)
            print("Get ssh_client by ssh pass.", flush=True)
        else:
            print("Can not connect host without ssh_pass or ssh_private_key!", flush=True)
            sys.exit(1)
    except Exception as err:
        print("连接服务器失败！", flush=True)
        print(err, flush=True)
        exit(1)
    return ssh_client


def checkEnv(sshClient):
    """检查部署主机环境是否符合要求"""
    cmd = driver_check_cmd
    stdin, stdout, stderr = sshClient.exec_command(cmd)
    inspectData = ""
    while True:
        rd = stdout.readline()
        inspectData = inspectData + rd
        if not rd:
            break
    err = stderr.read().decode("utf-8")
    if err:
        print(
            "nvidia驱动未安装或状态异常，请在目标主机执行`{0}`检查驱动情况，驱动安装下载请参考https://www.nvidia.cn/Download/index.aspx".format(
                driver_check_cmd), flush=True)
        sys.exit(1)


def chatglm_inspect(sshClient):
    """根据容器名检查是否存在及返回信息"""
    cmd = "{0} inspect {1}".format(docker_path, container_name)
    stdin, stdout, stderr = sshClient.exec_command(cmd)
    inspectData = ""
    while True:
        rd = stdout.readline()
        inspectData = inspectData + rd
        if not rd:
            break
    err = stderr.read().decode("utf-8")
    if err:
        print(err, flush=True)
        return None
    return json.loads(inspectData)


def sameConfig(config):
    """判断当前chatglm容器参数与需要部署的是否相同"""
    envList = config['Env']
    envDict = {}
    for env in envList:
        key = env.split("=")[0]
        value = "=".join(env.split("=")[1:])
        envDict[key] = value
    if glm_conc_cnt != envDict['CONC_CNT']:
        return False
    if glm_web_port != envDict['WEB_PORT']:
        return False
    if quantize_level != envDict['QUANTIZE_LEVEL']:
        return False
    if chatglm_image != config['Image']:
        return False
    return True


def delete_old_container(sshClient, containerName):
    """删除旧chatglm容器"""
    print("删除旧容器: {0}".format(containerName), flush=True)
    cmd_deleteContainer = "{0} stop {1} && {0} rm {1}".format(docker_path, containerName)
    print(cmd_deleteContainer, flush=True)
    stdin, stdout, stderr = sshClient.exec_command(cmd_deleteContainer)
    while True:
        rd = stdout.readline()
        print(rd, end="", flush=True)
        if not rd:
            break
    err = stderr.read().decode("utf-8")
    if err:
        print(err, flush=True)


def deploy_chatglm(sshClient):
    """部署chatglm容器"""
    print("创建容器: {0}".format(container_name), flush=True)
    deploy_cmd = "{0} run --gpus all -d --runtime=nvidia ".format(docker_path)
    deploy_cmd = deploy_cmd + "-e CONC_CNT={0} -e WEB_PORT={1} -e QUANTIZE_LEVEL={2} -p {1}:{1} ".format(glm_conc_cnt,
                                                                                                         glm_web_port,
                                                                                                         quantize_level)
    deploy_cmd = deploy_cmd + "--name={0} ".format(container_name)
    deploy_cmd = "{0} pull {1} && ".format(docker_path, chatglm_image) + deploy_cmd
    deploy_cmd = deploy_cmd + chatglm_image
    print(deploy_cmd, flush=True)
    stdin, stdout, stderr = sshClient.exec_command(deploy_cmd)
    while True:
        rd = stdout.readline()
        print(rd, end="", flush=True)
        if not rd:
            break
    err = stderr.read().decode("utf-8")
    if err:
        print(err, flush=True)
        if "Unknown runtime" in err:
            print(
                "请确保docker版本>19.03并已安装nvidia容器运行时，安装步骤参考https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/install-guide.html#setting-up-nvidia-container-toolkit",
                flush=True)
        sys.exit(1)


def check_inspect(sshClient, inspectInfo):
    """检查chatglm容器配置"""
    if not inspectInfo or len(inspectInfo) < 1:
        return True
    needRebuild = False
    containerInfo = inspectInfo[0]
    if containerInfo['State']['Status'] != "running" or not containerInfo['State']['Running']:
        needRebuild = True
        print("chatGLM container is not running.", flush=True)
    elif not sameConfig(containerInfo['Config']):
        needRebuild = True
        print("chatGLM config is changed.", flush=True)
    if needRebuild:
        delete_old_container(sshClient, containerInfo['Name'])
        return True
    print("已存在相同实例,不执行任何动作", flush=True)
    return False


def get_container_logs(sshClient, lines=100):
    """获取容器日志"""
    cmd_get_log = f"docker logs {container_name} --tail {lines}"
    stdin, stdout, stderr = sshClient.exec_command(cmd_get_log)
    logs = ""
    while True:
        rd = stdout.readline()
        logs = logs + rd
        if not rd:
            break
    err = stderr.read().decode("utf-8")
    if err:
        print("获取容器日志失败！", flush=True)
        print(err, flush=True)
    return logs


def check_inspect_health(sshClient, timeout_seconds):
    start_time = time.time()
    while True:
        time.sleep(30)
        if time.time() - start_time >= timeout_seconds:
            print("检测超时！服务启动失败，请前往服务器查看详情！", flush=True)
            logs = get_container_logs(sshClient)
            print(logs, flush=True)
            return False
        try:
            inspectInfo = chatglm_inspect(sshClient)
            if not inspectInfo or len(inspectInfo) < 1:
                print("Container {0} does not exist!".format(container_name), flush=True)
                continue
            elif inspectInfo[0]['State']['Health']['Status'] != "healthy":
                print("Health: {0}".format(inspectInfo[0]['State']['Health']['Status']), flush=True)
                logs = get_container_logs(sshClient)
                print(logs, flush=True)
                continue
            print("Health: {0}".format(inspectInfo[0]['State']['Health']['Status']), flush=True)
            return True
        except Exception as e:
            print(e, flush=True)
            continue


if __name__ == '__main__':
    sshClient = getSSHClient()
    checkEnv(sshClient)
    containerInfo = chatglm_inspect(sshClient)
    if check_inspect(sshClient, containerInfo):
        deploy_chatglm(sshClient)
    check_health = check_inspect_health(sshClient, timeout_seconds)
    if check_health:
        print("服务已成功启动！", flush=True)
    else:
        sys.exit(1)
